import json
import os
from copy import deepcopy

import numpy as np
import pandas as pd


def calculate_metric(df_file: str, top_num: int):
    df = pd.read_csv(df_file)
    all_gp = []
    rankic_list = []
    for code_market, gp in df.groupby("time"):
        oneday_data = deepcopy(gp)

        sorted_indices = np.argsort(-oneday_data["pred"])
        ranks = np.empty_like(sorted_indices)
        ranks[sorted_indices] = np.arange(len(sorted_indices))
        oneday_data["pred_rank"] = ranks
        oneday_data = oneday_data[oneday_data["pred_rank"] < top_num]

        sorted_indices = np.argsort(-oneday_data["change_ratio"])
        ranks = np.empty_like(sorted_indices)
        ranks[sorted_indices] = np.arange(len(sorted_indices))
        oneday_data["change_ratio_rank"] = ranks

        all_gp.append(oneday_data)
        tmp_rankic = oneday_data["change_ratio_rank"].corr(oneday_data["pred_rank"])
        rankic_list.append(tmp_rankic)
    # all_gp = pd.concat(all_gp, axis=1)
    # all_gp = all_gp[all_gp["pred_rank"] < top_num]

    rankic = np.array(rankic_list)
    meanrankic = np.nanmean(rankic)
    icir = np.nanmean(rankic) / np.nanstd(rankic)
    basedir = os.path.dirname(df_file)
    info = {"rankic": meanrankic, "icir": icir}
    with open(os.path.join(basedir, "icir_info.json"), "w") as f:
        json.dump(info, f, indent=2)
    print(json.dumps(info, indent=2))


if __name__ == "__main__":
    import sys

    pred_file = sys.argv[1]
    calculate_metric(pred_file, 10000)
